use memmap::{Mmap, MmapOptions};
use std::{
    fs::{File, OpenOptions},
    io::Write,
    sync::LazyLock,
};

static IMG: LazyLock<FatImage> = LazyLock::new(|| {
    let file = File::open("./abc.img").unwrap();
    let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };
    let ptr = mmap.as_ptr();
    unsafe {
        FatImage {
            bytes_per_sec: ptr.add(11).cast::<u16>().read_unaligned() as usize,
            sec_per_clus: ptr.add(13).cast::<u8>().read_unaligned() as usize,
            reserved_sec_count: ptr.add(14).cast::<u16>().read_unaligned() as usize,
            num_fats: ptr.add(16).cast::<u8>().read_unaligned() as usize,
            fat_size: ptr.add(36).cast::<u32>().read_unaligned() as usize,
            mmap,
        }
    }
});

#[derive(Debug)]
struct FatImage {
    bytes_per_sec: usize,
    sec_per_clus: usize,
    reserved_sec_count: usize,
    num_fats: usize,
    fat_size: usize,
    mmap: Mmap,
}

unsafe impl Send for FatImage {}
unsafe impl Sync for FatImage {}

trait ClusterOffset {
    fn cluster_offset(self) -> usize;
}

impl ClusterOffset for usize {
    fn cluster_offset(self) -> usize {
        let fat_region = IMG.num_fats * IMG.fat_size;
        IMG.bytes_per_sec * ((self - 2) * IMG.sec_per_clus + IMG.reserved_sec_count + fat_region)
    }
}

fn search_fat_img() {
    let img_size = 1024 * 1024 * 64;
    assert_eq!(img_size, IMG.mmap.len());
    println!("{:#?}", IMG);
    unsafe { search_directory(2) };
}

unsafe fn search_directory(cluster: usize) {
    let offset = cluster.cluster_offset();
    let base_ptr = IMG.mmap.as_ptr().add(offset);
    let mut entry_offset = 0;

    while entry_offset < IMG.sec_per_clus * IMG.bytes_per_sec {
        let entry_ptr = base_ptr.add(entry_offset);
        let entry = entry_ptr.cast::<[u8; 32]>().read_unaligned();

        if entry[0] == 0x00 {
            break;
        }
        if entry[0] == 0xE5 {
            entry_offset += 32;
            continue;
        }

        let is_dir = entry[11] == 0x10;
        let is_special = entry[0] == 0x2E; // . or ..
        let is_long_name = entry[11] == 0xf;

        if !is_special && !is_long_name && is_dir {
            let dir_cluster = dir_entry_cluster_num(&entry);
            search_directory(dir_cluster as usize);
        }
        if !is_long_name {
            print_entry_name(&entry);
        }

        entry_offset += 32;
    }

    let next_cluster = find_next_cluster(cluster);
    if next_cluster < 0x0FFFFFF8 {
        search_directory(next_cluster as usize);
    }
}

fn search_one_directory(base_ptr: *const u8) {
    unsafe {
        let mut entry_offset = 0;

        while entry_offset < IMG.sec_per_clus * IMG.bytes_per_sec {
            let entry_ptr = base_ptr.add(entry_offset);
            let entry = entry_ptr.cast::<[u8; 32]>().read_unaligned();

            if entry[0] == 0x00 {
                break;
            }
            if entry[0] == 0xE5 {
                entry_offset += 32;
                continue;
            }

            let is_dir = entry[11] == 0x10;
            let is_special = entry[0] == 0x2E; // . or ..
            let is_long_name = entry[11] == 0xf;

            if !is_long_name {
                print_entry_name(&entry);
            }

            entry_offset += 32;
        }
    }
}

fn parse_entry_name(entry: &[u8; 32]) -> String {
    let mut name = String::with_capacity(11);
    for &c in &entry[0..11] {
        if c.is_ascii_alphanumeric() {
            name.push(c as char);
        }
    }
    name
}

unsafe fn find_next_cluster(cluster: usize) -> u32 {
    let fat_offset = IMG.reserved_sec_count * IMG.bytes_per_sec;
    let fat_ptr = IMG.mmap.as_ptr().add(fat_offset);
    fat_ptr.add(cluster * 4).cast::<u32>().read_unaligned() & 0x0FFFFFFF
}

fn dir_entry_cluster_num(entry: &[u8; 32]) -> u32 {
    unsafe {
        let entry_ptr = entry.as_ptr();
        let cluster_hi = entry_ptr.add(20).cast::<u16>().read_unaligned();
        let cluster_lo = entry_ptr.add(26).cast::<u16>().read_unaligned();
        ((cluster_hi as u32) << 16) | cluster_lo as u32
    }
}

fn print_entry_name(entry: &[u8; 32]) {
    let name = parse_entry_name(entry);
    if !name.is_empty() {
        println!("{}", name);
    }
}

// 使用枚举明确分类类型
#[derive(Debug)]
enum ClusterType {
    Unused,
    DirectoryEntry,
    BmpHeader,
    BmpData,
}
fn main() {
    let mut nouse = Vec::new();
    let mut dir = Vec::new();
    let mut bmp_head = Vec::new();
    let mut bmp_data = Vec::new();
    collect_clusters(&mut nouse, &mut dir, &mut bmp_head, &mut bmp_data);

    for i in dir {
        let mut utility = ClusterUtility::new(i, cluster_bytes() as u32);
        while let Some(e) = utility.next() {
            println!(
                "{:?} is dir {} cluster num {}",
                e.name, e.is_dir, e.cluster_num
            );
            let bmp_head1 = cluster_num_to_ptr(e.cluster_num);
            let bmphead_resolv = bmp_head_resolv(bmp_head1);
            let mut recover = OpenOptions::new()
                .create(true)
                .write(true)
                .open(format!("./recover/{}", e.name))
                .unwrap();
            let mut buffer = Vec::with_capacity(bmphead_resolv.file_size as usize);
            unsafe {
                std::ptr::copy_nonoverlapping(
                    bmp_head1,
                    buffer.as_mut_ptr(),
                    bmphead_resolv.file_size as usize,
                );
                buffer.set_len(bmphead_resolv.file_size as usize);
            }
            recover.write_all(&buffer).unwrap();
        }
    }

    // let bmp_resolv = bmp_head_resolv(bmp_head[0]);
    // println!("{}", bmp_resolv.data.len());

    // unsafe { search_directory(2) };
}

fn cluster_num_to_ptr(cluster_num: u32) -> *const u8 {
    let offset = (cluster_num as usize).cluster_offset();
    unsafe { IMG.mmap.as_ptr().offset(offset as isize) }
}

struct ClusterUtility {
    ptr: *const u8,
    bytes: u32,
    offset: u32,
}
impl ClusterUtility {
    fn new(base_ptr: *const u8, bytes_per_cluster: u32) -> Self {
        Self {
            ptr: base_ptr,
            bytes: bytes_per_cluster,
            offset: 0,
        }
    }
}
impl Iterator for ClusterUtility {
    type Item = DirEntry;

    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            // 确保不会越界访问
            while self.offset + 32 <= self.bytes {
                let ptr = self.ptr.offset(self.offset as isize);
                let entry = ptr.cast::<[u8; 32]>().read_unaligned();

                match entry[0] {
                    0x00 => {
                        self.offset += 32;
                        continue;
                    } // 空条目
                    0xE5 => {
                        self.offset += 32;
                        continue;
                    } // 删除的条目
                    0x2E => {
                        self.offset += 32;
                        continue;
                    } // . 或 ..
                    _ => {}
                }

                let attr = entry[11];
                if attr == 0x0F {
                    // 长文件名条目
                    self.offset += 32;
                    continue;
                }

                // 验证是有效目录项
                if !(attr == 0x10 || attr == 0x20) {
                    self.offset += 32;
                    continue;
                }

                let cluster_num = dir_entry_cluster_num(&entry);
                let name = parse_entry_name(&entry);
                self.offset += 32;

                return Some(DirEntry {
                    ptr,
                    cluster_num,
                    name,
                    is_dir: attr == 0x10,
                });
            }
        }
        None
    }
}

fn collect_clusters(
    nouse_clus: &mut Vec<*const u8>,
    dir_clus: &mut Vec<*const u8>,
    bmp_head_clus: &mut Vec<*const u8>,
    bmp_data_clus: &mut Vec<*const u8>,
) {
    const ENTRY_SIZE: usize = 32;
    const BMP_SIGNATURE: [u8; 2] = [0x42, 0x4D]; // "BM"

    // 预计算簇步长避免循环中重复计算
    let cluster_step = (IMG.sec_per_clus * IMG.bytes_per_sec) as usize;
    // 使用with_capacity预分配空间提高性能
    let mut cluster_ptrs = Vec::with_capacity(IMG.mmap.len() / cluster_step);

    unsafe {
        let mut offset = 0x02_usize.cluster_offset();
        let ptr = IMG.mmap.as_ptr();

        while offset < IMG.mmap.len() {
            cluster_ptrs.push(ptr.add(offset));
            offset += cluster_step;
        }
    }

    let mut counts = [0; 4];
    // 对应四种分类的计数器

    cluster_ptrs.into_iter().for_each(|p| unsafe {
        let entry = p.cast::<[u8; ENTRY_SIZE]>().read_unaligned();

        let classification = match entry {
            [0, ..] => ClusterType::Unused,
            [0x2E, 0x20, ..] => ClusterType::DirectoryEntry, // 特殊目录项
            [0x2E, 0x2E, 0x20, ..] => ClusterType::DirectoryEntry, // 特殊目录项
            _ if entry[11] == 0x10
                || entry[11] == 0x20
                || entry[11] == 0x0F && entry[12] == 0x00 =>
            {
                ClusterType::DirectoryEntry
            }
            _ if entry.starts_with(&BMP_SIGNATURE) => ClusterType::BmpHeader,
            _ => ClusterType::BmpData,
        };

        match classification {
            ClusterType::Unused => {
                nouse_clus.push(p);
                counts[0] += 1;
            }
            ClusterType::DirectoryEntry => {
                dir_clus.push(p);
                counts[1] += 1;
            }
            ClusterType::BmpHeader => {
                bmp_head_clus.push(p);
                counts[2] += 1
            }
            ClusterType::BmpData => {
                bmp_data_clus.push(p);
                counts[3] += 1
            }
        }
    });

    println!("Unused clusters: {}", counts[0]);
    println!("Directory clusters: {}", counts[1]);
    println!("BMP headers clustes: {}", counts[2]);
    println!("BMP data clusters: {}", counts[3]);
}

fn bmp_head_resolv(bmp_head: *const u8) -> BMPHeadResolv {
    unsafe {
        let file_size = bmp_head.offset(0x2).cast::<u32>().read_unaligned();
        let data_offset = bmp_head.offset(0xa).cast::<u32>().read_unaligned();
        println!("data_offset {}", data_offset);
        let width = bmp_head.offset(0x12).cast::<u32>().read_unaligned();
        println!("width {}", width);
        let height = bmp_head.offset(0x16).cast::<u32>().read_unaligned();
        println!("height {}", height);
        let bits_per_pixl = bmp_head.offset(0x1c).cast::<u16>().read_unaligned();
        println!("bits_per_pixl {}", bits_per_pixl);
        let scan_line_bytes = (width * bits_per_pixl as u32 + 31) / 32 * 4;
        println!("scan_line_bytes {}", scan_line_bytes);

        let size = scan_line_bytes * height + 0x36 + 256 * 4;
        println!("cal size {}", size);

        println!("bmp file size {}", file_size);
        let data = if data_offset > cluster_bytes() as u32 {
            Vec::new()
        } else {
            let data_size = cluster_bytes() - data_offset as usize;
            let mut data = Vec::with_capacity(data_size);
            std::ptr::copy_nonoverlapping(
                bmp_head.offset(data_offset as isize),
                data.as_mut_ptr(),
                data_size,
            );
            data.set_len(data_size);
            data
        };

        BMPHeadResolv {
            file_size,
            data_offset,
            width,
            height,
            bits_per_pixl,
            scan_line_bytes,
            data,
        }
    }
}

fn cluster_bytes() -> usize {
    IMG.sec_per_clus * IMG.bytes_per_sec
}

struct BMPHeadResolv {
    file_size: u32,
    data_offset: u32,
    width: u32,
    height: u32,
    bits_per_pixl: u16,
    scan_line_bytes: u32,
    data: Vec<u8>,
}
#[derive(Debug)]
struct DirEntry {
    ptr: *const u8,
    cluster_num: u32,
    name: String,
    is_dir: bool,
}
